#!/usr/bin/env python3
"""
医学图像分割训练脚本

使用方法:
    python train.py
    python train.py --config custom_config.yaml
    python train.py --resume  # 从检查点恢复训练
"""

import argparse
import yaml
import os
import sys

from trainer import Trainer
from utils.model_utils import create_model_save_dirs
from data.data_loader import get_num_classes


def parse_args():
    """解析命令行参数"""
    
    parser = argparse.ArgumentParser(description='医学图像分割训练脚本')
    parser.add_argument('--config', type=str, default='config.yaml',
                       help='配置文件路径 (默认: config.yaml)')
    parser.add_argument('--resume', action='store_true',
                       help='从最新检查点恢复训练')
    parser.add_argument('--gpu', type=int, default=None,
                       help='指定GPU设备ID (默认: 自动选择)')
    
    return parser.parse_args()


def load_config(config_path):
    """加载配置文件"""
    
    if not os.path.exists(config_path):
        print(f"错误: 配置文件不存在: {config_path}")
        sys.exit(1)
    
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        print(f"成功加载配置文件: {config_path}")
        return config
    except Exception as e:
        print(f"错误: 无法加载配置文件: {e}")
        sys.exit(1)


def validate_config(config):
    """验证配置文件的完整性"""
    
    required_keys = [
        'data', 'model', 'training', 'augmentations'
    ]
    
    for key in required_keys:
        if key not in config:
            print(f"错误: 配置文件缺少必需的键: {key}")
            sys.exit(1)
    
    # 验证数据路径
    data_paths = [
        config['data']['image_dir'],
        config['data']['label_dir'],
        config['data']['train_list'],
        config['data']['val_list']
    ]
    
    for path in data_paths:
        if not os.path.exists(path):
            print(f"警告: 数据路径不存在: {path}")
    
    print("配置文件验证通过")


def setup_environment(args, config):
    """设置训练环境"""
    
    # 设置GPU
    if args.gpu is not None:
        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
        print(f"设置使用GPU: {args.gpu}")
    
    # 设置恢复训练
    if args.resume:
        config['training']['resume_checkpoint'] = True
        print("启用断点续训模式")
    
    # 创建必要的目录
    create_model_save_dirs(config)
    
    # 创建结果目录
    os.makedirs(os.path.dirname(config['evaluation']['metric_save_path']), exist_ok=True)
    os.makedirs(config['inference']['output_dir'], exist_ok=True)


def print_config_summary(config):
    """打印配置摘要"""
    
    print("\n" + "="*50)
    print("训练配置摘要")
    print("="*50)
    print(f"模型: {config['model']['name']}")
    print(f"类别数: {get_num_classes(config)}")
    print(f"图像尺寸: {config['data']['image_size']}")
    print(f"批次大小: {config['training']['batch_size']}")
    print(f"学习率: {config['training']['learning_rate']}")
    print(f"训练轮数: {config['training']['epochs']}")
    print(f"优化器: {config['training']['optimizer']}")
    print(f"损失函数: {config['training']['loss_function']}")
    print(f"数据增强: {'启用' if config['augmentations']['enabled'] else '禁用'}")
    if config['augmentations']['enabled']:
        print(f"增强停止轮数: {config['augmentations']['augment_stop_epoch']}")
    print(f"模型保存目录: {config['training']['save_dir']}")
    print(f"日志目录: {config['training']['log_dir']}")
    print("="*50 + "\n")


def main():
    """主函数"""
    
    # 解析命令行参数
    args = parse_args()
    
    # 加载配置
    config = load_config(args.config)
    
    # 验证配置
    validate_config(config)
    
    # 设置环境
    setup_environment(args, config)
    
    # 打印配置摘要
    print_config_summary(config)
    
    try:
        # 创建训练器并开始训练
        trainer = Trainer(config)
        model = trainer.train()
        
        print("\n" + "="*50)
        print("训练完成！")
        print("="*50)
        print(f"最佳模型保存在: {os.path.join(config['training']['save_dir'], 'best_model.pth')}")
        print(f"训练日志保存在: {config['training']['log_dir']}")
        print(f"训练曲线保存在: {os.path.join(config['training']['log_dir'], 'training_curves.png')}")
        
    except KeyboardInterrupt:
        print("\n训练被用户中断")
        sys.exit(0)
    except Exception as e:
        print(f"\n训练过程中发生错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()

